import torch
from sklearn.metrics import accuracy_score
from models.emotion_lstm import EmotionLSTM
from utils.dataset_loader import load_dataset
from config import *
import matplotlib.pyplot as plt

def evaluate():
    # 加载数据
    X_train, X_test, y_train, y_test = load_dataset(DATA_PATH, max_pad_len=MAX_PAD_LEN)

    # 转换为PyTorch张量
    X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
    y_test_tensor = torch.tensor(y_test, dtype=torch.long)

    # 初始化模型
    model = EmotionLSTM(INPUT_DIM, HIDDEN_DIM, OUTPUT_DIM, NUM_LAYERS)
    model.load_state_dict(torch.load(MODEL_PATH))

    # 测试模型
    model.eval()  # 设置为评估模式

    # 用于绘图的记录变量
    test_accuracies = []

    with torch.no_grad():
        outputs = model(X_test_tensor)
        _, predicted = torch.max(outputs, 1)
        accuracy = accuracy_score(y_test, predicted.numpy())
        test_accuracies.append(accuracy)

    print(f'Test Accuracy: {accuracy:.4f}')

    # 绘制测试准确率曲线（如果有多轮评估或不同模型）
    plt.figure(figsize=(6, 5))
    plt.plot(test_accuracies, label='Test Accuracy', color='green')
    plt.title('Test Accuracy Curve')
    plt.xlabel('Test Evaluations')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.show()

if __name__ == '__main__':
    evaluate()
